import sys
from pathlib import Path
import random
from datetime import datetime, timedelta

# 将项目根目录添加到 Python 路径
ROOT_DIR = Path(__file__).parent.parent
sys.path.append(str(ROOT_DIR))

from app import create_app, db
from app.models.pathogen import Pathogen

def create_test_data():
    """创建测试数据"""
    app = create_app()
    
    with app.app_context():
        # 病原体名称列表
        virus_names = ['H5N1', 'H7N9', 'H9N2', 'NDV', 'IBV']
        bacteria_names = ['大肠杆菌', '沙门氏菌', '支原体', '布鲁氏菌', '巴氏杆菌']
        parasite_names = ['球虫', '蓝舌病毒', '弓形虫', '蛔虫', '绦虫']
        fungus_names = ['曲霉', '毛霉', '青霉', '白色念珠菌', '隐球菌']
        
        # 地点列表
        locations = [
            ('新疆阿勒泰', 47.8, 88.1),
            ('新疆伊犁', 43.9, 81.3),
            ('新疆巴州', 41.8, 86.1),
            ('新疆喀什', 39.5, 75.9),
            ('新疆阿克苏', 41.2, 80.3)
        ]
        
        # 为每种类型创建测试数据
        pathogen_types = {
            'virus': virus_names,
            'bacteria': bacteria_names,
            'parasite': parasite_names,
            'fungus': fungus_names
        }
        
        test_data = []
        for type_name, names in pathogen_types.items():
            for name in names:
                # 随机选择一个地点
                location = random.choice(locations)
                # 随机生成日期（过去一年内）
                isolation_date = datetime.now() - timedelta(days=random.randint(1, 365))
                
                pathogen = Pathogen(
                    name=name,
                    type=type_name,
                    isolation_date=isolation_date,
                    isolation_location=location[0],
                    latitude=location[1],
                    longitude=location[2],
                    gene_sequence='ATCG' * random.randint(10, 20),
                    created_by=1  # 假设是管理员创建的
                )
                test_data.append(pathogen)
        
        # 批量添加数据
        db.session.bulk_save_objects(test_data)
        db.session.commit()
        
        print(f'成功添加 {len(test_data)} 条测试数据！')

if __name__ == '__main__':
    print("开始添加测试数据...")
    create_test_data()
    print("添加完成！") 